#include <iostream>

extern "C" void set_cuda_device(int device_id) {
  cudaError_t err = cudaSetDevice(device_id);
  if (err != cudaSuccess) {
    std::cerr << "Failed to set device: " << cudaGetErrorString(err)
              << std::endl;
    return;
  }
}

extern "C" void get_device_count(int *count) { cudaGetDeviceCount(count); }
